import numpy as np
import onnxruntime as rt

from utils.image_utils import resize_image
from Detector.data_process import resizedImage2YoloInput, output_process, confidence_filter, NMS, re_letterbox

# 暂存区，之后写到配置文件里
INPUT_SIZE = (640, 640)
CONF_THRES = 0.3
IOU_THRES = 0.7


class Detector(object):
    def __init__(self, onnx_path):
        self._load_onnx(onnx_path=onnx_path)

    def __call__(self, image: np.ndarray):
        return self._run(image)

    def _load_onnx(self, onnx_path):
        self.sess = rt.InferenceSession(onnx_path)
        self.input_name = self.sess.get_inputs()[0].name
        self.label_name = self.sess.get_outputs()[0].name

    def _run(self, image: np.ndarray):
        resized_image = resize_image(image, new_size=INPUT_SIZE)
        input_tensor = resizedImage2YoloInput(resized_image)
        output_tensor = self.sess.run(
            [self.label_name],
            {self.input_name: input_tensor}
        )[0]
        output_tensor = output_process(output_tensor)
        output_tensor = confidence_filter(output_tensor, conf_threshold=CONF_THRES)
        output_tensor = NMS(output_tensor, IoU_threshold=IOU_THRES)
        output = re_letterbox(output_tensor, image)  # xyxy, conf, class
        return output


if __name__ == '__main__':
    import cv2 as cv
    from config.paths import DETECTOR, TEST_IMAGE_PATH

    img = cv.imread(TEST_IMAGE_PATH)
    detector = Detector(onnx_path=DETECTOR['ONNX_PATH'])
    print(detector(img))
